import re
import sys
sys.path.append("/Users/ender/Code_repo/rosalind-brush-questions")
from utils import read_seqfile


def parser_codon_table(file, stop_symbol = "*"):
    codon_table = {}
    with open(file, "r") as f:
        tmp_table_file = "      ".join([line.strip("\n") for line in f.readlines()])
        for codon_aa in re.findall("\S{3}\s\S{1,4}", tmp_table_file):
            c, aa = codon_aa.split(" ")
            if aa == "Stop":
                codon_table[c] = stop_symbol
            else:
                codon_table[c] = aa

    return codon_table

def rna2protein(rna: str, codon_table: dict) -> str:
    protein = ""
    for i in range(0, len(rna), 3):
        protein += codon_table[rna[i:i+3]]
    return protein


rna = read_seqfile("./rosalind_prot.txt")
codon_table = parser_codon_table("./codon_table.txt")
protein = rna2protein(rna, codon_table)
print(protein)